module flash_rw (
    input  logic clk,
    input  logic rst_n,
    input  logic start,
    // Add your ports here
    output logic led,
    output logic done
);

  //========================================================
  // Parameter and Localparam Declarations
  //========================================================
  localparam TEST_VAL = 8'hAB;
  localparam TEST_ADDR = 8'h01;
  //========================================================
  // Signal Declarations
  //========================================================
  // Add your internal signals here

  logic [7:0] data_in = 8'd0;
  logic [7:0] data_out;
  logic spi_start = 1'b0;
  logic spi_busy;
  logic sclk;
  logic mosi;
  logic miso;
  logic cs_n;
  logic flash_cs_n;
  logic [7:0] val = TEST_VAL;
  logic [1:0] bytes_cnt = 0;

  typedef enum logic [1:0] {
    IDLE = 3'b000,
    WRITE_ADDR = 3'b001,
    WRITE_DATA = 3'b010,
    READ_ADDR = 3'b011,
    READ_DATA = 3'b100,
    DONE = 3'b101
  } state_t;

  state_t current_state, next_state;
  always_ff @(posedge clk or negedge rst_n) begin
    if (!rst_n) begin
      current_state <= IDLE;
    end else begin
      current_state <= next_state;
    end
  end
  //========================================================
  // Combinational Logic
  //========================================================
  always_comb begin
    if (!rst_n) begin
      next_state = IDLE;
    end else begin
      case (current_state)
        IDLE: begin
          if (start) begin
            next_state = WRITE_ADDR;
          end else begin
            next_state = IDLE;
          end
        end
        WRITE_ADDR: begin
          if (spi_busy == 0) begin
            next_state = WRITE_DATA;  // Transition to WRITE_DATA after WRITE_ADDR
          end else begin
            next_state = WRITE_ADDR;
          end
        end
        READ: begin
          // Add your read logic here
          if (spi_busy == 0 && val == 8'hff) begin
            next_state = DONE;  // Transition to WRITE after READ
          end else begin
            next_state = READ;
          end
          // next_state = WRITE;
        end
        WRITE: begin
          if (spi_busy == 0 && val == 8'h00) begin
            next_state = READ;  // Transition to WRITE after READ
          end else begin
            next_state = WRITE;
          end
        end
        DONE: begin
          next_state = IDLE;
        end
        default: begin
          next_state = IDLE;
        end
      endcase
    end
  end

  //========================================================
  // Sequential Logic
  //========================================================
  always_ff @(posedge clk or negedge rst_n) begin
    if (!rst_n) begin
      spi_start <= 1'b0;
      data_in <= 8'd0;
      flash_cs_n <= 1'b1;  // Active low
      bytes_cnt <= 0;
      done <= 1'b0;
    end else if (current_state == DONE) begin
      done <= 1'b1;  // Indicate done
    end else if (spi_start == 1 && spi_busy == 1) begin
      spi_start <= 1'b0;
    end else if (current_state == IDLE && start) begin
      data_in <= val;
      val <= TEST_VAL;  // Reset or hold the data_in as needed
      bytes_cnt <= 0;
      flash_cs_n <= 1'b1;  // Active low
      if (data_out == TEST_VAL) begin
        led <= 1'b1;  // Indicate success
      end else begin
        led <= 1'b0;  // Indicate failure
      end
    end else if (current_state == WRITE && spi_busy == 0) begin
      if (bytes_cnt == 0) begin
        data_in   <= TEST_ADDR;
        bytes_cnt <= 1;
      end else begin
        data_in <= val;
        val <= 8'hff;
      end
      flash_cs_n <= 1'b0;  // Active low
      spi_start  <= 1'b1;
    end else if (current_state == READ && spi_busy == 0) begin
      if (bytes_cnt == 0) begin
        data_in   <= TEST_ADDR;
        bytes_cnt <= 1;
      end else begin
        val <= 8'd0;  // Reset or hold the data_in as needed
      end
      spi_start  <= 1'b1;
      flash_cs_n <= 1'b0;  // Active low
    end else begin
      flash_cs_n <= 1'b1;  // Active low
      done <= 1'b0;
    end
  end

  //========================================================
  // Instantiations
  //========================================================
  spi_master #(
      .CPOL(0),
      .CPHA(0)
  ) u_spi_master (
      .clk     (clk),
      .rst_n   (rst_n),
      .data_in (data_in),
      .start   (spi_start),
      .sclk    (sclk),
      .mosi    (mosi),
      .miso    (miso),
      .cs_n    (cs_n),
      .busy    (spi_busy),
      .data_out(data_out)
  );

endmodule
